# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of Cybertron package.
#
# The Cybertron is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Graph embedding module for molecular representation
"""

import torch
from torch import nn, Tensor
from typing import Union, Tuple, Optional

from ..activation import get_activation
from ..units import Units, GLOBAL_UNITS
from ..utils import get_integer, GLOBAL_DEVICE, get_arguments

_EMBEDDING_BY_KEY = {}

# Define the _embedding_register function
def _embedding_register(*aliases):
    """Return the alias register."""
    def alias_reg(cls):
        name = cls.__name__.lower()
        if name not in _EMBEDDING_BY_KEY:
            _EMBEDDING_BY_KEY[name] = cls

        for alias in aliases:
            if alias not in _EMBEDDING_BY_KEY:
                _EMBEDDING_BY_KEY[alias] = cls
        return cls
    return alias_reg


# Define the GraphEmbedding class
class GraphEmbedding(nn.Module):
    r"""Base class for graph embedding modules.

    Args:
        dim_node (int): Dimension of node embedding vectors.
        dim_edge (int): Dimension of edge embedding vectors.
        emb_dis (bool or int, optional): Whether to include distance information in embedding.
            Default: True
        emb_bond (bool or int, optional): Whether to include bond information in embedding.
            Default: False
        activation (Any, optional): Activation function to use.
            Default: None
        length_unit (str, optional): Unit for length measurements.
            Default: None
    """
    def __init__(self,
                 dim_node: int,
                 dim_edge: int,
                 emb_dis: Union[bool, int] = True,
                 emb_bond: Union[bool, int] = False,
                 activation: Union[nn.Module, str] = None,
                 length_unit: Optional[str] = GLOBAL_UNITS.length_unit,
                 **kwargs
                 ):
        super().__init__()
        self._kwargs = get_arguments(locals(), kwargs)

        if length_unit is None:
            length_unit = GLOBAL_UNITS.length_unit
        self.units = Units(length_unit)

        self.emb_dis = get_integer(emb_dis)
        self.emb_bond = get_integer(emb_bond)

        self._dim_node = get_integer(dim_node)
        self._dim_edge = get_integer(dim_edge)

        self.cutoff = None

        self.activation = get_activation(activation)
        self.device = GLOBAL_DEVICE()

    @property
    def dim_node(self):
        return self._dim_node

    @property
    def dim_edge(self):
        return self._dim_edge

    def _print_info(self,
                    num_retraction: int = 3,
                    num_gap: int = 3,
                    char: str = ' '
                    ) -> None:
        """print the information of molecular model"""
        ret = char * num_retraction
        gap = char * num_gap
        print(ret + f' Graph Embedding: {self._get_name()}')
        print('-' * 80)
        print(ret + gap + f' Length unit: {self.units.length_unit}')
        print(ret + gap + f' Embedding distance: {self.emb_dis}')
        print(ret + gap + f' Embedding Bond: {self.emb_bond}')
        print(ret + gap + f' Dimension of node embedding vector: {self.dim_node}')
        print(ret + gap + f' Dimension of edge embedding vector: {self.dim_edge}')
        print('-' * 80)

    def convert_length_from(self, unit: Union[str, Units]) -> float:
        """returns a scale factor that converts the length from a specified unit."""
        return self.units.convert_length_from(unit)

    def convert_length_to(self, unit: Union[str, Units]) -> float:
        """returns a scale factor that converts the length to a specified unit."""
        return self.units.convert_length_to(unit)

    def forward(self,
                atom_type: Tensor,
                atom_mask: Tensor,
                distance: Tensor,
                dis_mask: Tensor,
                bond: Tensor,
                bond_mask: Tensor,
                **kwargs,
                ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
        """Compute the properties of the molecules.

        Args:
            atom_type (Tensor): Tensor of shape (B, A). Data type is int.
                Index of atom types. Default: None
            atom_mask (Tensor): Tensor of shape (B, A). Data type is bool
                Mask for atom types
            distance (Tensor): Tensor of shape (B, A, A). Data type is float.
                Distances between central atom and its neighouring atoms.
            dis_mask (Tensor): Tensor of shape (B, A, A). Data type is bool.
                Mask for neighbour list.
            bond (Tensor): Tensor of shape (B, A, A). Data type is int.
                Types index of bond connected with two atoms
            bond_mask (Tensor): Tensor of shape (B, A, A). Data type is bool.
                Mask for bonds

        Returns:
            node_emb (Tensor): Tensor of shape (B, A, E). Data type is float.
                Node embedding vector.
            node_mask (Tensor): Tensor of shape (B, A, E). Data type is float.
                Mask for Node embedding vector.
            edge_emb (Tensor): Tensor of shape (B, A, A, K). Data type is float.
                Edge embedding vector.
            edge_mask (Tensor): Tensor of shape (B, A, A, K). Data type is float.
                Mask for edge embedding vector.
            edge_cutoff (Tensor): Tensor of shape (B, A, A). Data type is float.
                Cutoff for edge.

        Symbols:
            B:  Batch size.
            A:  Number of atoms in system.
            E:  Dimension of node embedding vector
            K:  Dimension of edge embedding vector
            D:  Spatial dimension of the simulation system. Usually is 3.

        """
        raise NotImplementedError
